package stars.enums;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

public class StarDistributions {
	
	private HashMap<String, Double> youngProbabilities = new HashMap<String, Double>();
	private HashMap<String, Double> averageProbabilities = new HashMap<String, Double>();
	private HashMap<String, Double> oldProbabilities = new HashMap<String, Double>();
	private ArrayList<String> entries = new ArrayList<String>();
	
	public ArrayList<String> getEntries() {
		return this.entries;
	}
	
	public void addProbability(String param, double value, String age) {
		if (age.equalsIgnoreCase("young")) {
			youngProbabilities.put(param, value);
		}
		else if (age.equalsIgnoreCase("average")) {
			averageProbabilities.put(param, value);
		}
		else {
			oldProbabilities.put(param, value);
		}
		
		if (!entries.contains(param)) {
			entries.add(param);
		}
	}
	
	public double getProbability(String param, String age) {
		if (age.equalsIgnoreCase("young")) {
			return youngProbabilities.get(param);
		}
		else if (age.equalsIgnoreCase("average")) {
			return averageProbabilities.get(param);
		}
		else 
			return oldProbabilities.get(param);
	}
	
	public HashMap<String, Double> getAllProbabilities(String age) {
		if (age.equalsIgnoreCase("young")) {
			return youngProbabilities;
		}
		else if (age.equalsIgnoreCase("average")) {
			return averageProbabilities;
		}
		else 
			return oldProbabilities;
	}
	
	public void setProbabilities(HashMap<String, Double> parameters, String age) {
		if (age.equalsIgnoreCase("young")) {
			this.youngProbabilities.clear();
			this.youngProbabilities.putAll(parameters);
		}
		else if (age.equalsIgnoreCase("average")) {
			this.averageProbabilities.clear();
			this.averageProbabilities.putAll(parameters);
		}
		else { 
			this.oldProbabilities.clear();
			this.oldProbabilities.putAll(parameters);
		}
		
		for (Map.Entry<String, Double> entry : parameters.entrySet()) {
			if (!entries.contains(entry.getKey())) {
				entries.add(entry.getKey());
			}
		}
	}
	
	public void normalize() {
		
		double cumuProbYoung = 0.0;
		double cumuProbAverage = 0.0;
		double cumuProbOld = 0.0;		
		
		for (String entry: entries) {
			cumuProbYoung += youngProbabilities.get(entry);
			cumuProbAverage += averageProbabilities.get(entry);
			cumuProbOld += oldProbabilities.get(entry);			
		}
		
		double youngCorrection = 1;
		double averageCorrection = 1;
		double oldCorrection = 1;
		
		if (cumuProbYoung != 1.0)
			youngCorrection = 1/cumuProbYoung;
		
		if (cumuProbAverage!=1.0)
			averageCorrection = 1/cumuProbAverage;
		
		if (cumuProbOld!=1.0) 
			oldCorrection = 1/cumuProbOld;
		
		double correctedYoung;
		double correctedAverage;
		double correctedOld;
		
		for (String entry: entries) {
			correctedYoung = youngProbabilities.get(entry)*youngCorrection;
			correctedAverage = averageProbabilities.get(entry)*averageCorrection;
			correctedOld = oldProbabilities.get(entry)*oldCorrection;
			
			youngProbabilities.put(entry, correctedYoung);
			averageProbabilities.put(entry, correctedAverage);
			oldProbabilities.put(entry, correctedOld);
		}
	}

}
